import mne
import numpy as np
import os
import glob
import h5py
import warnings
from sklearn.preprocessing import StandardScaler
import datetime  # 用于获取当前时间，以便在日志中加上时间戳


# EEG数据文件路径
eeg_data_path = 'G:\\Research\\EEG_Project\\sleepstages_CNN\\data\\eegdata'
eeg_anno_path = 'G:\\Research\\EEG_Project\\sleepstages_CNN\\data\\annotation'
preprocessed_dir = 'G:\\Research\\EEG_Project\\sleepstages_CNN\\data\\data_anno_savefiles'
Gauss_preprocessed_dir = 'G:\\Research\\EEG_Project\\sleepstages_CNN\\data\\data_anno_savefiles_Gauss'
file_names=[]
#文件存储路径 如果有哪种新的处理方式需要修改存放路径 将相应变量置为1即可
Gauss_Flag= 1     #决定数据保存文件的路径


#从注释中获取未知标签描述，此前默认未知描述为-1
def get_unknown_descriptions(annotations, unknown_indices, segment_length, sfreq):
    """
    从annotations对象和未知标签的索引中获取未知描述。
    """
    # 转换为对应的原始时间点
    unknown_times = unknown_indices * segment_length / sfreq
    unknown_desc = set()
    # 获取开始时间和结束时间内包含的注释
    for time in unknown_times:
        for annot in annotations:
            if annot['onset'] <= time < annot['onset'] + annot['duration']:
                unknown_desc.add(annot['description'])
    return unknown_desc
#向日志中写入记录
def log_message(message):
    """
    向指定的日志文件写入消息。
    """
    # 创建日志文件的完整路径
    if Gauss_Flag:
        log_file_path = os.path.join(Gauss_preprocessed_dir, '数据预处理日志.txt')
    else:
        log_file_path = os.path.join(preprocessed_dir, '数据预处理日志.txt')
    # 获取当前时间并格式化
    now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    with open(log_file_path, 'a') as log_file:  # 以追加模式打开文件
        log_file.write(f"{now} - {message}\n")  # 写入时间戳和消息

#高斯噪声增强
def add_gaussian_noise(data, noise_level=0.01):
    noise = np.random.normal(0, noise_level, data.shape)
    data_noisy = data + noise
    return data_noisy

#数据标准化
def standardize_data(data):
    # MNE已经支持多种方式归一化数据
    # 这里使用了简单的z-score标准化方法
    scaler = StandardScaler()
    standardized_data = scaler.fit_transform(data.reshape(-1, data.shape[-1])).reshape(data.shape)
    return standardized_data

#滤波
def apply_bandpass_filter(data, sfreq, l_freq=0.5, h_freq=30):
    # 初始化被移除片段的索引列表
    removed_segments = []
    # 对每个片段的数据进行带通滤波
    for segment in range(data.shape[0]):
        try:
            # MNE的滤波函数需要二维数组
            segment_data = data[segment, :].reshape((1, -1))
            # 应用带通滤波
            filtered_data = mne.filter.filter_data(segment_data, sfreq, l_freq, h_freq, verbose=False)
            # 将滤波后的数据放回原数组
            data[segment, :] = filtered_data
        except Exception as e:
            # 滤波失败时，记录这个片段的索引，并打印警告信息
            warnings.warn(f'滤波时出现问题：{e}')
            removed_segments.append(segment)
    # 如果有片段因滤波失败而被移除，更新数据和段数
    if removed_segments:
        data = np.delete(data, removed_segments, axis=0)
        # 返回被移除片段的索引列表
        return data, removed_segments
    return data, None

def Data_prepare():
    print("数据预处理")
    global Guass
    edf_names = glob.glob(os.path.join(eeg_data_path, '*.edf'))
    #样本总数
    number_samples = len(edf_names)
    #创建文件名元组  供主函数打开使用
    global file_names
    #保存相关文件
    for i in range(number_samples):
        ########################初次读取文件数据##########################
        name = os.path.basename(edf_names[i])
        name = name[:-9] #去掉后缀
        #对应的注释文件名
        hyp_name = glob.glob(os.path.join(eeg_anno_path, name + '*' + '-Hypnogram.edf'))
        hyp_name = ''.join(hyp_name) #转字符串
        # 数据保存文件
        if Gauss_Flag:
            save_file = os.path.join(Gauss_preprocessed_dir, name + '.h5')
        else:
            save_file = os.path.join(preprocessed_dir, name + '.h5')
        # 将保存文件存储 便于查看 此变量已设置为全局变量
        file_names += (save_file,)

        if not os.path.exists(save_file):
            # 使用mne读取EDF文件
            print("正在处理第%d个文件" %i)
            log_message(f"\r\n\##################文件{name}处理日志:##################")
            raw_data = mne.io.read_raw_edf(edf_names[i], preload=True)
            annotations = mne.read_annotations(hyp_name)
            #log_message(f"原始数据长度:{len(raw_data)} 原始注释长度：{len(annotations)}")

            ########################原始数据片段化##########################
            #print(raw_data.info)
            #print(annotations.description)
            # 选择特定通道的数据
            picked_channel = 'EEG Fpz-Cz'
            data, times = raw_data[picked_channel, :]
            #将每个片段设置为30秒的数据。首先，获取采样频率
            sfreq = int(raw_data.info['sfreq'])
            # 计算每30秒片段包含的数据点数
            segment_length = 30 * sfreq
            # 确定将所有数据分割成30秒片段所需的片段数
            n_segments = int(data.shape[1] // segment_length)
            # 为简化起见，如果数据不能被完整地划分成30秒的片段，这里忽略最后一部分
            data = data[:, :n_segments * segment_length]# 保留所有行 但是只保留n_segments * segment_length列
            # 将数据重塑成三维数组 (n_segments, segment_length, 1)
            # 对应于 (样本数, 时间步长, 通道数)
            data = data.reshape((n_segments, segment_length, 1))
            #log_message(f"未经处理的片段数:{n_segments}")

            ########################添加相应注释,并且对应数量##########################
            stages_mapping = {
                'Sleep stage W':  0,
                'Sleep stage 1':  1,
                'Sleep stage 2':  2,
                'Sleep stage 3':  3,
                'Sleep stage 4':  3,  # 经常将阶段3和4合并为N3
                'Sleep stage R':  4
            }
            # 初始化阶段标签数组
            stage_labels = np.zeros(n_segments)

            # 遍历注释，为每一个30秒的EEG片段标记正确的睡眠阶段
            for annot in annotations:
                start_idx = int(annot['onset'] * sfreq) // segment_length  # 找到注释开始的索引
                end_idx = int((annot['onset'] + annot['duration']) * sfreq) // segment_length  # 找到注释结束的索引
                label = stages_mapping.get(annot['description'], -1)  # 获取阶段标签，默认为-1（未知阶段)
                stage_labels[start_idx:end_idx] = label  # 设置标签
            # 计算并记录原始数据的各类标签数量
            unique_labels, counts = np.unique(stage_labels, return_counts=True)
            label_count_message = "原始数据的标签类别数量: " + ", ".join(
                f"{label}: {count} " for label, count in zip(unique_labels, counts))
            log_message(label_count_message)

            if np.any(stage_labels == -1):
                warnings.warn("有未知着阶段标签。")
                # 获取所有未知标签的索引
                unknown_label_indices = np.where(stage_labels == -1)[0]
                # 获取未知标签的描述
                unknown_descriptions = get_unknown_descriptions(annotations, unknown_label_indices, segment_length,sfreq)
                # 记录未知标签的描述
                log_message(f"文件{name}: 存在未知的睡眠阶段标签，数量: {len(unknown_label_indices)}，描述: {unknown_descriptions}")
                # 移除这些未知标签
                stage_labels = np.delete(stage_labels, unknown_label_indices)
                # 移除对应的EEG数据
                data = np.delete(data, unknown_label_indices, axis=0)
                # 更新段数，因为可能有数据被移除
                n_segments = len(stage_labels)
            # 对于不完整(过短)的最后一个时间段，这里删除它
            stage_labels = stage_labels[:n_segments]
            log_message(f"处理后的片段数以及标签数(未滤波)：{n_segments,len(stage_labels)}")

            ########################数据剔除##########################
            # 首先找出哪个标签类别是最多的
            unique_labels, label_counts = np.unique(stage_labels, return_counts=True)
            max_count_label = unique_labels[np.argmax(label_counts)]
            # 计算标签数量最多的索引
            max_label_indices = np.where(stage_labels == max_count_label)[0]
            # 确定要保留的标签为0的数据比例，这里是0.3
            keep_ratio = 0.3
            # 计算要删除的标签数量
            num_to_delete = int(len(max_label_indices) * (1 - keep_ratio))
            # 随机选择要删除的索引
            indices_to_delete = np.random.choice(max_label_indices, num_to_delete, replace=False)
            # 创建一个布尔数组，初始为True，之后将要删除的索引设置为False
            keep_indices = np.ones(len(data), dtype=bool)
            keep_indices[indices_to_delete] = False
            # 过滤掉要删除的数据
            data = data[keep_indices]
            stage_labels = stage_labels[keep_indices]
            # 更新段数，因为数据已经变化
            n_segments = len(data)
            log_message(f"删除最多标签后的片段数和标签数: {n_segments}, {len(stage_labels)}")

            ########################滤波处理##########################
            # 将数据重塑成二维数组 (n_segments, segment_length), 因为MNE滤波函数需要二维数组
            data = data.reshape((n_segments, segment_length))
            # 使用滤波函数进行带通滤波
            data, removed_segments = apply_bandpass_filter(data, sfreq)
            # 更新标签，若有需要
            if removed_segments is not None:
                stage_labels = np.delete(stage_labels, removed_segments)
                n_segments = data.shape[0]
                log_message(f"滤波后的数据长度:{n_segments}")
            # 确保滤波后的数据和标签数量相符
            assert len(stage_labels) == n_segments, "滤波后，标签和样本数量不一致。"
            # 将滤波后的数据重塑回三维数组 (样本数, 时间步长, 通道数) 以适应深度学习模型输入需求
            data = data.reshape((n_segments, segment_length, 1))

            ########################数据标准化并确认标签数是否正常##########################
            # 数据标准化
            data = standardize_data(data)
            unique_labels = np.unique(stage_labels)
            if 4 <= len(unique_labels) <= 5:
                print(f"标签检查通过，找到的唯一标签数量在4到5之间，具体为: {len(unique_labels)}")
                #for label in unique_labels:
                    # 这里假设有一个反向映射从标签值映射回标签名称，如果没有则直接打印label值
                    #print(f"- 标签值{label}: {list(stages_mapping.keys())[list(stages_mapping.values()).index(label)]}")
            else:
                print("警告: 标签的唯一数量异常。")
                print(f"找到的标签有: {unique_labels}")
                input("按任意键继续...")

            ##############################数据增强##############################
            if Gauss_Flag:
                # 在这里，计算各个睡眠阶段的样本数，并进行数据增强
                # 获得了每个睡眠阶段的样本数
                unique, counts = np.unique(stage_labels, return_counts=True)
                max_count = np.max(counts)
                # 创建用于存储增强后数据的列表
                data_augmented = data.tolist()
                stage_labels_augmented = stage_labels.tolist()
                # 数据增强：针对样本数少于max_count的类别进行增强
                for stage, count in zip(unique, counts):
                    if count < max_count:
                        # 计算需要增强的样本数
                        num_augmented = max_count - count
                        # 获取当前阶段的样本
                        stage_data = data[stage_labels == stage]
                        # 循环增强样本，直到达到max_count
                        for _ in range(num_augmented):
                            # 随机选择一个样本进行增强
                            sample_to_augment = stage_data[np.random.randint(stage_data.shape[0])]
                            # 添加高斯噪声
                            augmented_sample = add_gaussian_noise(sample_to_augment)
                            # 添加到数据中
                            data_augmented.append(augmented_sample.reshape(segment_length, 1))
                            stage_labels_augmented.append(stage)
                # 将增强数据转换为numpy数组
                data_augmented = np.array(data_augmented)
                stage_labels_augmented = np.array(stage_labels_augmented)
                log_message(f"最终增强数据和标签长度: {len(data_augmented)}, {len(stage_labels_augmented)}")
            else:
                # 验证结果
                assert len(stage_labels)==n_segments, "标签和样本不一致"
                log_message(f"最终数据和标签长度:{n_segments},{len(stage_labels)}")

            ##############################数据文件保存##############################
            #保存数据文件
            # 创建一个新的HDF5文件
            with h5py.File(save_file, 'w') as f:
                # 在文件中创建一个新的组
                grp = f.create_group('eeg')
                # 在这个组中创建两个新的数据集，分别保存 `eeg_data` 和 `stage_labels`
                if Gauss_Flag:
                    grp.create_dataset('eeg_data', data=data_augmented)
                    grp.create_dataset('stage_labels', data=stage_labels_augmented)
                    # 计算并记录增强后的各类标签数量
                    unique_labels_aug, counts_aug = np.unique(stage_labels_augmented, return_counts=True)
                    label_count_message_aug = "增强后的标签类别数量: " + ", ".join(
                        f"{label}: {count}" for label, count in zip(unique_labels_aug, counts_aug))
                    log_message(label_count_message_aug)
                else:
                    grp.create_dataset('eeg_data', data=data)
                    grp.create_dataset('stage_labels', data=stage_labels)
                    # 计算并记录原始数据的各类标签数量
                    unique_labels, counts = np.unique(stage_labels, return_counts=True)
                    label_count_message = "原始数据的标签类别数量: " + ", ".join(
                        f"{label}: {count}" for label, count in zip(unique_labels, counts))
                    log_message(label_count_message)
        else:
            print("第%d个文件%s 已经存在" % (i, save_file))


